import numpy as np
import polars as pl


def generate_ev_available(days: int) -> pl.DataFrame:
    np.random.seed(0)
    arrival_times = np.random.choice(
        [18, 19, 20, 21, 22], size=(days,), p=[0.30, 0.30, 0.20, 0.15, 0.05]
    )
    depature_times = np.random.choice([6, 7, 8], size=(days,), p=[0.2, 0.4, 0.4])
    arr = np.zeros(shape=(days, 24))
    for d in range(days):
        arr[d, arrival_times[d] :] = 1
        arr[d, : depature_times[d]] = 1
    return pl.from_numpy(arr, schema={f"{h}:00": pl.UInt8 for h in range(24)})


def generate_ev_reduction(days: int) -> pl.DataFrame:
    np.random.seed(0)
    energy = np.random.beta(2, 3, size=(days,)) * (20 - 5) + 5
    return pl.DataFrame({"ev_level_reduction": energy})


generate_ev_available(93).write_csv("envs/assets/ev_available.csv")
generate_ev_reduction(93).write_csv("envs/assets/ev_reduction.csv")
